stan_code <- "
data {
  
  // number of country-year observations in each age group
  int N_obs;                              

  int cat1_start;                    // start index for iterations associated with latitude >40
  int cat1_end;                      // end index for iterations associated with latitude >40
  int cat2_start;                    // start index for iterations associated with latitude 20-40
  int cat2_end;                      // end index for iterations associated with latitude 20-40
  int cat3_start;                    // start index for iterations associated with latitude <20
  int cat3_end;                      // end index for iterations associated with latitude <20

  real pop_04[N_obs];                // population for 0-4 age group 
  real pop_514[N_obs];               // population for 5-14 age group
  real p_tx_it[N_obs];               // CDR for both age groups
  real hiv_04[N_obs];                // HIV prevalence for 0-4 age group
  real hiv_514[N_obs];               // HIV prevalence for 5-14 age group
  real pem_04[N_obs];                // PEM prevalence for 0-4 age group
  real pem_514[N_obs];               // PEM prevalence for 5-14 age group
  real bcg_04[N_obs];                // BCG coverage for 0-4 age group
  real bcg_514[N_obs];               // BCG coverage for 5-14 age group
  real art[N_obs];                   // pediatric ART coverage
  real contacts_sum_04[N_obs];       // sum of contacts for 0-4 age group
  real contacts_sum_514[N_obs];      // sum of contacts for 5-14 age group
  real contacts_inc_04[N_obs];       // dot product of adult incidence rates and contacts for 0-4 age group  
  real contacts_inc_514[N_obs];      // dot product of adult incidence rates and contacts for 5-14 age group
  int who_cases_04[N_obs];           // WHO TB case notifications for 0-4 age group
  int who_cases_514[N_obs];          // WHO TB case notifications for 5-14 age group

  real mn[3];                        // rate ratios dependent on latitude 

  // prior distribution parameters
  real infect_a;
  real infect_b;
  real active_04_a;
  real active_04_b;
  real active_514_a;
  real active_514_b;
  real dur_untx_a;
  real dur_untx_b;
  real fract_early_a;
  real fract_early_b;
  real inv_sqrt_phi_sd;
  real hiv_rr_a;
  real hiv_rr_b;
  real pem_rr_a;
  real pem_rr_b;
  real art_rr_a;
  real art_rr_b;
  real rr_ped_detect_a;  //~~
  real rr_ped_detect_b;  //~~
  real shape_rate;

} 

///////////////////////////////////////////////////////////
transformed data {
  // dont need anything here right now
  
}
///////////////////////////////////////////////////////////
parameters {
  
  real<lower=0, upper=1> infect;               // probability of infection per contact
  real<lower=0, upper=1> active_04;            // probability of active disease for 0-4 age group
  real<lower=0, upper=1> active_514;           // probability of active disease for 5-14 age group
  real<lower=0> dur_untx;                      // duration of untreated disease
  real<lower=0,upper=1> fract_early;           // fraction of time spent in early disease
  real<lower=0>  inv_sqrt_phi;                 // inverse square root of negative binomial dispersion factor
  real<lower=0, upper=1> sat;                  // magnitude of contact saturation effect in households 
  real<lower=0> hiv_rr;                        // HIV risk ratio
  real<lower=0> pem_rr;                        // PEM risk ratio
  real<lower=0> art_rr;                        // ART risk ratio
  real<lower=0> rr_ped_detect;                 // relative probability of detection for ped vs adults, in high performing settings  //~~
  real<lower=0> z;                             // parameter controlling BCG risk ratios
  
} 
///////////////////////////////////////////////////////////
transformed parameters {
  
  real duration;                               // calculated duration value
  real phi;                                    // negative binomial dispersion factor
  real bcg_rr_cat1;                            // BCG risk ratio associated with latitude >40
  real bcg_rr_cat2;                            // BCG risk ratio associated with latitude 20-40
  real bcg_rr_cat3;                            // BCG risk ratio associated with latitude <20
  
  vector[N_obs] notif_04;                      // notification estimates for 0-4 age group
  vector[N_obs] notif_514;                     // notification estimates for 5-14 age group

  phi = pow(inv_sqrt_phi, -2);
  
  bcg_rr_cat1 = z * mn[1];
  bcg_rr_cat2 = z * mn[2];
  bcg_rr_cat3 = z * mn[3];
  
  // model for pediatric TB incidence
  for(i in cat1_start:cat1_end){
    
    duration = dur_untx * (1 - p_tx_it[i] * (1 - fract_early));  //~~
    
    notif_04[i]  = p_tx_it[i] * rr_ped_detect * pop_04[i] * active_04  //~~
                              * (1-hiv_04[i] + hiv_04[i]*(1-art[i])* hiv_rr + hiv_04[i]*art[i]*hiv_rr*art_rr)
                              * (1 - pem_04[i] + pem_rr      * pem_04[i])
                              * (1 - bcg_04[i] + bcg_rr_cat1 * bcg_04[i])
                              * (contacts_inc_04[i]/contacts_sum_04[i])
                              * ((1 - exp(-infect * duration
                              * contacts_sum_04[i] * sat)) / sat);
    
    notif_514[i] = p_tx_it[i] * rr_ped_detect * pop_514[i] * active_514  //~~
                              * (1-hiv_514[i] + hiv_514[i]*(1-art[i])* hiv_rr + hiv_514[i]*art[i]*hiv_rr*art_rr)
                              * (1 - pem_514[i] + pem_rr      * pem_514[i])
                              * (1 - bcg_514[i] + bcg_rr_cat1 * bcg_514[i])
                              * (contacts_inc_514[i]/contacts_sum_514[i])
                              * ((1 - exp(-infect * duration
                              * contacts_sum_514[i] * sat)) / sat);

  }
  
  for(i in cat2_start:cat2_end){
    
    duration = dur_untx * (1 - p_tx_it[i] * (1 - fract_early));  //~~
    
    notif_04[i]  = p_tx_it[i] * rr_ped_detect * pop_04[i] * active_04  //~~
                              * (1-hiv_04[i] + hiv_04[i]*(1-art[i])* hiv_rr + hiv_04[i]*art[i]*hiv_rr*art_rr)
                              * (1 - pem_04[i] + pem_rr      * pem_04[i])
                              * (1 - bcg_04[i] + bcg_rr_cat2 * bcg_04[i])
                              * (contacts_inc_04[i]/contacts_sum_04[i])
                              * ((1 - exp(-infect * duration
                              * contacts_sum_04[i] * sat)) / sat);
    
    notif_514[i] = p_tx_it[i] * rr_ped_detect * pop_514[i] * active_514  //~~
                              * (1-hiv_514[i] + hiv_514[i]*(1-art[i])* hiv_rr + hiv_514[i]*art[i]*hiv_rr*art_rr)
                              * (1 - pem_514[i] + pem_rr      * pem_514[i])
                              * (1 - bcg_514[i] + bcg_rr_cat2 * bcg_514[i])
                              * (contacts_inc_514[i]/contacts_sum_514[i])
                              * ((1 - exp(-infect * duration
                              * contacts_sum_514[i] * sat)) / sat);

  }
  
  for(i in cat3_start:cat3_end){
    
    duration = dur_untx * (1 - p_tx_it[i] * (1 - fract_early));  //~~
    
    notif_04[i]  = p_tx_it[i] * rr_ped_detect * pop_04[i] * active_04  //~~
                              * (1-hiv_04[i] + hiv_04[i]*(1-art[i])* hiv_rr + hiv_04[i]*art[i]*hiv_rr*art_rr)
                              * (1 - pem_04[i] + pem_rr      * pem_04[i])
                              * (1 - bcg_04[i] + bcg_rr_cat3 * bcg_04[i])
                              * (contacts_inc_04[i]/contacts_sum_04[i])
                              * ((1 - exp(-infect * duration
                              * contacts_sum_04[i] * sat)) / sat);
    
    notif_514[i] = p_tx_it[i] * rr_ped_detect * pop_514[i] * active_514  //~~
                              * (1-hiv_514[i] + hiv_514[i]*(1-art[i])* hiv_rr + hiv_514[i]*art[i]*hiv_rr*art_rr)
                              * (1 - pem_514[i] + pem_rr      * pem_514[i])
                              * (1 - bcg_514[i] + bcg_rr_cat3 * bcg_514[i])
                              * (contacts_inc_514[i]/contacts_sum_514[i])
                              * ((1 - exp(-infect * duration
                              * contacts_sum_514[i] * sat)) / sat);

  }  
  
}
///////////////////////////////////////////////////////////
model {
  
  // priors
  infect ~ beta(infect_a, infect_b);
  active_04 ~ beta(active_04_a, active_04_b);
  active_514 ~ beta(active_514_a, active_514_b);
  dur_untx  ~ gamma(dur_untx_a , dur_untx_b);
  fract_early  ~ beta(fract_early_a , fract_early_b);
  inv_sqrt_phi  ~ normal(0, inv_sqrt_phi_sd);
  sat ~ beta(2, 2);
  hiv_rr ~ gamma(hiv_rr_a, hiv_rr_b);
  pem_rr ~ gamma(pem_rr_a, pem_rr_b);
  art_rr ~ gamma(art_rr_a, art_rr_b);
  rr_ped_detect ~ beta(rr_ped_detect_a, rr_ped_detect_b); //~~
  z ~ gamma(shape_rate, shape_rate);
  // likelihood
  who_cases_04 ~ neg_binomial_2(notif_04, phi);
  who_cases_514 ~ neg_binomial_2(notif_514, phi);

} 
///////////////////////////////////////////////////////////
generated quantities {
  // dont need anything here right now

}
"
